from remote_pb2 import WriteRequest
from types_pb2 import  Label, Sample, TimeSeries, MetricMetadata
from snappy import compress
import requests
import time

# Remote Write 端点
remote_write_url = "http://localhost:9529/prom_remote_write"

metric_name = "baihongru_custom_metric"

# 构建 MetricMetadata 对象
metadata = MetricMetadata(
    type = 1,  # 指标类型，为 GAUGE 类型
    metric_family_name = metric_name,  # 指标名称，与 TimeSeries 中的 __name__ 字段一致
    help = "This is a custom metric provided by baihongru.",
    unit = "ms",
)

# 创建 Labels
# 以 `__` 开始的字段是保留 Labels 似乎只有 __name__
# 字段定义见：https://github.com/prometheus/prometheus/blob/main/model/labels/labels_common.go
labels = [
    Label(
        name = "__name__",
        value = metric_name,
    ),
    Label(
        name = "foo",
        value = "bar",
    )
]

# 创建 Samples，包含值和时间戳
samples = [
    Sample(
        value = 42.0,
        timestamp = int(time.time() * 1000)  # 毫秒时间戳
    )
]

# 构建 TimeSeries 对象
time_series = TimeSeries(
    labels=labels,
    samples=samples
)

# 将 TimeSeries 对象添加到 WriteRequest 中
write_req = WriteRequest()
write_req.metadata.extend([metadata])
write_req.timeseries.extend([time_series])

# 设置请求头
headers = {
    "Content-Type": "application/x-protobuf",
    "Content-Encoding": "snappy",
}

# 压缩数据
serialized_data = write_req.SerializeToString()
compressed_data = compress(serialized_data)

try:
    # 发送 POST 请求到 Remote Write 端点
    response = requests.post(remote_write_url, data=compressed_data, headers=headers)
    if response.status_code == 204:
        print("Data successfully sent to the Remote Write endpoint.")
    else:
        print(f"Failed to send data. Status code: {response.status_code}, Reason: {response.text}")
except requests.RequestException as e:
    print(f"Error occurred while sending data: {str(e)}")